import csv
import random
from datetime import datetime, timedelta

import numpy as np
from scipy.stats import poisson

from osc02.static.constants import Constants

'''
We define the entire dataset rainfall_data as a dictionary where the keys represent dates ('MMDD') 
and the values are lists of 0s and 1s for each day. You can add the remaining 364 days of data in this
 dictionary.
We set a random seed for reproducibility to ensure consistent results.
We loop through each day's record in rainfall_data, calculate the average rainfall for the day, 
fit a Poisson distribution, and generate a random sample to simulate the daily rainfall count.
The simulated daily rainfall counts are stored in the simulated_rainfall dictionary.
Finally, we print the simulated daily rainfall counts with the respective date keys.
'''


# the return is a dict with
# the string keys representing desired period in 2012 (not specially for 2012 but to make up a full date)
# and the list values indicating precipitation times on the day over 12 years
def get_sampled_precipitation_times_in_period(seed=1,
                                              start_date=datetime(2023, 3, 10),
                                              end_date=datetime(2023, 7, 30)):
    # Define your entire dataset as a dictionary where keys represent dates ('MMDD') and
    # values are lists of 0s and 1s.

    precipitation_data = Constants.PRECIPITATION_TIMES_ON_DAYS_IN_12_YEARS
    # Set a random seed for reproducibility
    random.seed(seed)

    # Initialize an empty dictionary to store simulated daily precipitation counts
    simulated_precipitation_times_in_specified_days = {}

    precipitation_data_in_period_dict = \
        get_sliced_precipitation_data_in_period_dict(start_date, end_date, precipitation_data)
    with open('precipitation.csv', 'w', encoding='utf-8', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['日期', '最高温', '最低温', '天气', '风力动向'])
        # Loop through each day's record and sample precipitation
        for day_key, day_record in precipitation_data_in_period_dict.items():
            average_precipitation = np.mean(day_record)  # Calculate the average precipitation for the day
            poisson_fit = poisson(average_precipitation)  # Fit a Poisson distribution
            # return is a numPy array containing one Poisson random variable.
            sample = poisson_fit.rvs(size=1)  # Generate a random sample for the day
            simulated_precipitation_times_in_specified_days[day_key] = sample[0]
            writer.writerow([day_key, sample[0]])
    return simulated_precipitation_times_in_specified_days


# called by get_sampled_precipitation_times_in_period
# get sliced precipitation data in the period between start_date and end_date
def get_sliced_precipitation_data_in_period_dict(start_date, end_date, precipitation_data_dict):
    current_date = start_date
    day_delta = timedelta(days=1)
    sliced_precipitation_data_in_period_dict = {}
    while current_date <= end_date:
        key = Constants.YEAR_STR_REPRESENTING_PRECIPITATION \
              + f"{current_date.month:0{2}}" + f"{current_date.day:0{2}}"
        # copy desired element to precipitation_data_specified
        sliced_precipitation_data_in_period_dict[key] = precipitation_data_dict.get(key)
        current_date += day_delta
    return sliced_precipitation_data_in_period_dict


if __name__ == '__main__':
    print(get_sampled_precipitation_times_in_period())
